package com.study.feature.transform

import org.apache.spark.ml.feature.NGram
import org.apache.spark.sql.SparkSession

/**
 * 特征转换-NGram(得到组合序列)
 *
 * NGram是一个Transformer，用来将序列形式的数据(通常是单词序列)，按指定的整数N组合后成为新的序列。
 * 该算法通常被用在自然语言处理(nlp)任务的数据预处理步骤中。
 *
 * NGram模型本身是计算机语言学上的一种常见的统计概率模型，
 * 在语言处理中经常需要通过一组单词的序列来判断后续单词出现的概率，这时就需要使用n元组合的数据作为输入。
 *
 * @author stephen
 * @date 2019-08-28 13:48
 */
object NGramDemo {

  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    val wordDataFrame = spark.createDataFrame(Seq(
      (0, Array("Hi", "I", "heard", "about", "Spark")),
      (1, Array("I", "wish", "Java", "could", "use", "case", "classes")),
      (2, Array("Logistic", "regression", "models", "are", "neat"))
    )).toDF("id", "words")

    val ngram = new NGram()
      // 定义用多少个序列值做组合(默认值:2)
      .setN(2)
      .setInputCol("words")
      .setOutputCol("ngrams")

    val ngramDataFrame = ngram.transform(wordDataFrame)
    ngramDataFrame.show(false)

    //    +---+------------------------------------------+------------------------------------------------------------------+
    //    |id |words                                     |ngrams                                                            |
    //    +---+------------------------------------------+------------------------------------------------------------------+
    //    |0  |[Hi, I, heard, about, Spark]              |[Hi I, I heard, heard about, about Spark]                         |
    //    |1  |[I, wish, Java, could, use, case, classes]|[I wish, wish Java, Java could, could use, use case, case classes]|
    //    |2  |[Logistic, regression, models, are, neat] |[Logistic regression, regression models, models are, are neat]    |
    //    +---+------------------------------------------+------------------------------------------------------------------+
  }
}
